#!/usr/bin/env python3

import argparse
import logging
import re
import shutil
import subprocess as p
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, Sequence

log = logging.getLogger()

REMOVE_LOC = Path(__file__).parent / 'remove_loc.sh'


def install(*args, show_locals: bool = True, **kwargs):
    try:
        import rich
        from rich.logging import RichHandler
        from rich.traceback import install as tb_install

        logging.basicConfig(
            level='NOTSET',
            format='%(message)s',
            datefmt='[%X]',
            handlers=[
                RichHandler(),
            ],
        )
        global log
        log = logging.getLogger()

        global print
        print = rich.print

        old_handler = tb_install(
            *args,
            show_locals=show_locals,
            **kwargs,
        )
        return old_handler

    except ImportError:
        logging.basicConfig(
            level='NOTSET',
        )
        return None


class Options(argparse.Namespace):
    input: Path
    output: Path
    dump_dir: Path
    build_dir: Path
    replace: bool
    verbose: int

    @staticmethod
    def create_parser():
        p = argparse.ArgumentParser()
        p.add_argument(
            '--build-dir',
            type=Path,
            required=True,
            help='panda build dir',
        )
        p.add_argument(
            '--dump-dir',
            type=Path,
            required=True,
            help='directory for dumps',
        )
        p.add_argument(
            '-r',
            '--replace',
            help='Clear ets-dump-dir',
            action=argparse.BooleanOptionalAction,
            default=False,
        )
        p.add_argument(
            '-v',
            '--verbose',
            action='count',
            default=0,
        )
        p.add_argument(
            'input',
            type=Path,
            help='input files',
        )
        p.add_argument(
            'output',
            type=Path,
            help='output file',
        )
        return p

    @staticmethod
    def parse(parser: argparse.ArgumentParser | None = None):
        p = parser or Options.create_parser()
        return p.parse_args(namespace=Options())


class es2panda:
    def __init__(self, build_path: Path):
        self._build = build_path

    def arktsconfig(self):
        return self._build / 'tools' / 'es2panda' / 'generated' / 'arktsconfig.json'

    def es2panda(self):
        return (
            str((self._build / 'bin' / 'es2panda').expanduser()),
            '--extension=ets',
        )

    def phases(self) -> list[str]:
        cmd = [
            *self.es2panda(),
            '--list-phases',
        ]
        proc = p.run(cmd, stdout=p.PIPE, stderr=p.STDOUT, text=True, check=False)
        return re.findall(r'^ (\w+)$', proc.stdout, re.M)

    def help(self):
        cmd = [
            *self.es2panda(),
            '--help',
        ]
        proc = p.run(cmd, stdout=p.PIPE, stderr=p.STDOUT, text=True, check=False)
        return proc.stdout

    def build(self, *args: str, input: Path, output: Path):
        cmd = [
            *self.es2panda(),
            '--gen-stdlib=false',
            f'--arktsconfig={self.arktsconfig().expanduser()!s}',
            *args,
            f'--output={output!s}',
            f'{input!s}',
        ]
        proc = p.run(cmd, stdout=p.PIPE, stderr=p.STDOUT, text=True, check=False)
        return proc.stdout

    def build_with_dump_ets(self, input: Path, output: Path):
        phases = self.phases()
        log.debug('Phases: %s', phases)
        phases_str = ','.join(phases)

        return self.build(
            f'--dump-ets-src-after-phases={phases_str}',
            f'--dump-after-phases={phases_str}',
            input=input,
            output=output,
        )


_ETS_PATTERN = re.compile(r'After phase (\w+) ets source:')
_AST_PATTERM = re.compile(r'After phase (\w+):')


@dataclass
class File:
    name: str
    lines: list[str] = field(default_factory=list)


@dataclass
class Result:
    ets: list[File] = field(default_factory=list)
    ast: list[File] = field(default_factory=list)


def _ets_dump_parser(text: str):
    result = Result()
    lines = text.splitlines(keepends=True)

    def match(pattern: re.Pattern, line: str, files: list[File]):
        if (m := pattern.match(line)) is not None:
            f = File(name=str(m.group(1)))
            files.append(f)
            return f.lines

    cur: list[str] = []
    for line in lines:
        if (m := match(_ETS_PATTERN, line, result.ets)) is not None:
            cur = m
        elif (m := match(_AST_PATTERM, line, result.ast)) is not None:
            cur = m
        else:
            cur.append(line)
    return result


def _prepare_out_dir(dir: Path, clear: bool):
    if clear and dir.exists():
        shutil.rmtree(dir)
        log.info('%s has been cleared', dir.resolve())
    if not dir.exists():
        dir.mkdir(parents=True)
        log.info('ETS dump directory created %s', dir.resolve())


def _write_file(file: Path, lines: Iterable[str]):
    with file.open(mode='w', encoding='utf-8') as fp:
        fp.writelines(lines)


def _write_files(dir: Path, files: Iterable[File], first_id: int, suffix: str):
    _prepare_out_dir(dir, False)
    result: list[Path] = []
    for id, file in enumerate(files, first_id):
        f = dir / f'{id:02d}-{file.name}{suffix}'
        _write_file(f, file.lines)
        log.debug(str(f))
        result.append(f)
    return result


def _no_loc(dir: Path, json_files: Sequence[Path]):
    for file in json_files:
        p.run(f'"{REMOVE_LOC!s}" "{file!s}"', shell=True)


def main(options: Options):
    log.debug(options)
    panda = es2panda(options.build_dir)
    output = panda.build_with_dump_ets(
        input=options.input,
        output=options.output,
    )
    log.debug(str(options.output))
    result = _ets_dump_parser(output)
    result.ets = [
        File(
            name='Original',
            lines=options.input.read_text().splitlines(keepends=True),
        ),
        *result.ets,
    ]
    out_dir = options.dump_dir
    _prepare_out_dir(
        dir=out_dir,
        clear=options.replace,
    )
    _write_files(out_dir / 'ets', result.ets, 0, '.ets')
    json_files = _write_files(out_dir / 'ast', result.ast, 1, '.json')
    _no_loc(out_dir / 'ast', json_files)


def _log_level(verbose: int):
    log_levels = [
        logging.WARNING,
        logging.INFO,
        logging.DEBUG,
        logging.NOTSET,
    ]
    n = max(min(len(log_levels) - 1, verbose), 0)
    return log_levels[n]


if __name__ == '__main__':
    install()
    options = Options.parse()
    log.setLevel(level=_log_level(options.verbose))
    main(options)
